from bzrc import BZRC, Command, ShootCommand
import numpy, sys, math, time, optparse, random, pprint
from numpy import array, mat

TIME_STEP = .36
FRIC_COEFF = .005
F = mat([[1, TIME_STEP, TIME_STEP * TIME_STEP / 2, 0, 0, 0],
	[0, 1, TIME_STEP, 0, 0, 0],
	[0, -FRIC_COEFF, 1, 0, 0, 0],
	[0, 0, 0, 1, TIME_STEP, TIME_STEP * TIME_STEP / 2],
	[0, 0, 0, 0, 1, TIME_STEP],
	[0, 0, 0, 0, -FRIC_COEFF, 1]])
FT = F.T
H = mat([[1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0]])
HT = H.T
EX = mat([[.1, 0, 0, 0, 0, 0], [0, .1, 0, 0, 0, 0], [0, 0, 10, 0, 0, 0], [0, 0, 0, .1, 0, 0], [0, 0, 0, 0, .1, 0], [0, 0, 0, 0, 0, 10]])
MU = mat(numpy.zeros((6, 1), float))
SIGMA = mat([[10, 0, 0, 0, 0, 0], [0, .1, 0, 0, 0, 0], [0, 0, .1, 0, 0, 0], [0, 0, 0, 10, 0, 0], [0, 0, 0, 0, .1, 0], [0, 0, 0, 0, 0, .1]])
I = mat(numpy.identity(6, float))
EZ = []
SHOT_SPEED = 100 #shot speed is about 100 meters a second
RANGE = 345


class Agent(object):
	def __init__(self, bzrc):
		self.bzrc = bzrc
		self.constants = self.bzrc.get_constants()
		self.commands = []
		self.lastTime = time.time()
		self.plot = Plotter()

	def tick(self):
		'''Some time has passed; decide what to do next'''
		# Get information from the BZRC server
		mytanks = self.bzrc.get_mytanks()
		othertanks = self.bzrc.get_othertanks()
		# Reset my set of commands (we don't want to run old commands)
		self.commands = []

		# Decide what to do with my tank
		self.behave(mytanks[0], othertanks[0])
		
		# Send the commands to the server
		results = self.bzrc.do_commands(self.commands)
		
		if othertanks[0].status != "alive":
			return True
		return False
			
	def behave(self, bot, target):
		startTime = time.time()
		global MU
		global SIGMA
		v = mat([[(target.x - MU[0,0]) / TIME_STEP], [(target.y - MU[3,0]) / TIME_STEP]])
		a = mat([[(v[0,0] - MU[1,0]) / TIME_STEP, (v[1,0] - MU[4,0]) / TIME_STEP]]).T
		X = mat([(target.x, v[0,0], a[0,0], target.y, v[1,0], a[1,0])]).T
		Z = H * X
		FET = F * SIGMA * FT + EX
		K = FET * HT * (H * FET * HT + EZ).I
		FMU = F * MU
		MU = FMU + K * (Z - H * FMU)
		SIGMA = (I - K * H) * FET
		
		self.plot.plot(SIGMA[0,0], SIGMA[3,3], MU[0,0], MU[3,0])
		
		self.aim(bot, startTime, F)
		
	def aim(self, bot, startTime, F):
		global MU
		distToTarget = math.hypot(bot.x - MU[0,0], bot.y - MU[3,0])
			
		timeToTarget = distToTarget / SHOT_SPEED + .5
		
		FPrime = mat([[1, timeToTarget, timeToTarget * timeToTarget / 2, 0, 0, 0],
			[0, 1, timeToTarget, 0, 0, 0],
			[0, -FRIC_COEFF, 1, 0, 0, 0],
			[0, 0, 0, 1, timeToTarget, timeToTarget * timeToTarget / 2],
			[0, 0, 0, 0, 1, timeToTarget],
			[0, 0, 0, 0, -FRIC_COEFF, 1]]) #Matrix for extrapolating position to when the shot will arrive
			
		projectedLocationMU = FPrime * MU #FPrime * MU
		projectedLocation = (projectedLocationMU[0,0], projectedLocationMU[3,0])
		
		delAng = 10000
		now = time.time()
		
		#while (delAng < -.03 or delAng > .03) and (now - startTime + .07) < TIME_STEP:		
		desiredAngle = math.atan2(projectedLocation[1] - bot.y, projectedLocation[0] - bot.x)
		delAng = self.normalize_angle(desiredAngle - bot.angle)
		command = Command(bot.index, 0, delAng)
		self.commands.append(command)
		# Send the commands to the server
		results = self.bzrc.do_commands(self.commands)
		mytanks = self.bzrc.get_mytanks()
		bot = mytanks[bot.index]
		now = time.time()
		delAng = self.normalize_angle(desiredAngle - bot.angle)
		
		distToTarget = math.hypot(bot.x - projectedLocation[0], bot.y - projectedLocation[1])
		
		if(distToTarget > RANGE or math.fabs(delAng) > .04):
			return #Don't bother shooting at someone out of range or not in our line of sight
		command = ShootCommand(bot.index)
		self.commands.append(command)
			
	def normalize_angle(self, angle):
		'''Make any angle be between +/- pi.'''
		angle -= 2 * math.pi * int (angle / (2 * math.pi))
		if angle <= -math.pi:
			angle += 2 * math.pi
		elif angle > math.pi:
			angle -= 2 * math.pi
		return angle
		
class Plotter(object):
	def __init__(self):
		try:
			from Gnuplot import GnuplotProcess
		except ImportError:
			print "Sorry.  You don't have the Gnuplot module installed."
			import sys
			sys.exit(-1)
		self.gp = GnuplotProcess(persist=1)
	
	def gnuplot_header(self, minimum, maximum):
		'''Return a string that has all of the gnuplot sets and unsets.'''
		s = ''
		s += 'set xrange [%s: %s]\n' % (minimum, maximum)
		s += 'set yrange [%s: %s]\n' % (minimum, maximum)
		s += 'set pm3d\n'
		s += 'set view map\n'
		# The key is just clutter.  Get rid of it:
		s += 'unset key\n'
		# Make sure the figure is square since the world is square:
		s += 'set size square\n'
		# Add a pretty title (optional):
		s += "set title 'Search'\n"
		s += "set palette model RGB functions 1-gray, 1-gray, 1-gray\n"
		s += "set isosamples 100\n"
		return s
		
	def plot_Gaussian(self, sigx, sigy, x, y):
		s = 'sigma_x = %s\n \
			sigma_y = %s\n \
			rho = 0\n \
			splot 1.0/(2.0 * pi * sigma_x * sigma_y * sqrt(1 - rho**2)) \
			* exp(-1.0/2.0 * ((x - %s)**2 / sigma_x**2 + (y - %s)**2 / sigma_y**2 \
			- 2.0*rho*(x - %s)*(y - %s)/(sigma_x*sigma_y))) with pm3d\n' % (sigx, sigy, x, y, x, y)
		return s
	
	def plot(self, sigx, sigy, x, y):
		self.gp.write(self.gnuplot_header(-400, 400))
		self.gp.write(self.plot_Gaussian(sigx, sigy, x, y))
		self.gp.flush()
	

def main():
	# Process CLI arguments.
	parser = optparse.OptionParser(usage = "%prog [options]",
								   version = "%prog 0.1")
	parser.add_option("-o", "--host", type="string",dest="host",
			metavar="HOST",default='localhost',
			help="hostname for the server")
	parser.add_option("-p", "--port", type="int",dest="port",
			metavar="PORT",default=3907,
			help="port number for the server")
	parser.add_option("-c", "--covariance", type="int",dest="covariance",
			metavar="COVARIANCE",default=5,
			help="Covariance constant for this lab")

	(options, args) = parser.parse_args()

	# Connect.
	#bzrc = BZRC(options.host, int(options.port), debug=True)
	bzrc = BZRC(options.host, int(options.port))
	
	stdDev = options.covariance
	global EZ
	EZ = mat([[stdDev * stdDev, 0], [0, stdDev * stdDev]])

	agent = Agent(bzrc)

	prevTime = time.time()
	# Run the agent
	hit = False
	try:
		while not hit:
			hit = agent.tick()
			now = time.time()
			timeDiff = now - prevTime
			print "Took ", timeDiff, " time to calculate."
			if (TIME_STEP > timeDiff):
				time.sleep(TIME_STEP - timeDiff)
			prevTime = time.time()
			
	except KeyboardInterrupt:
		print "Exiting due to keyboard interrupt."
		bzrc.close()


if __name__ == '__main__':
	main()
